from PyQt5.QtWidgets import QWidget, QVBoxLayout, QLabel, QHBoxLayout, QFileDialog
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QFont
from qfluentwidgets import SmoothScrollArea, PrimaryPushButton, InfoBar, InfoBarPosition
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg
import os

# 定义颜色列表
COLOR_LIST = ['#4CC2FF', 'orange', 'green', 'purple', 'red', 'blue', 'yellow', 'pink']
current_color_index = 1


def read_cost_data(file_path, mode='drone'):
    """读取总成本和不同类型数据"""
    num_agents = []  # 无人机数量、障碍物数量或空域大小
    total_costs = []  # 总成本
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            for line in f:
                try:
                    if mode == 'drone':
                        agents, cost = map(float, line.strip().split(',')[:2])
                        num_agents.append(int(agents))
                        total_costs.append(cost)
                    elif mode == 'obstacle':
                        agents, cost, obstacle = map(float, line.strip().split(',')[:3])
                        num_agents.append(int(obstacle))
                        total_costs.append(cost)
                    elif mode == 'airspace':
                        agents, cost, obstacle, airspace_str = line.strip().split(',')
                        try:
                            num_agents.append(airspace_str)
                            total_costs.append(float(cost))
                        except (IndexError, ValueError):
                            continue
                except ValueError:
                    continue
    except Exception as e:
        print(f"读取数据出错: {e}")
        return [], []
    return num_agents, total_costs


class StatsInterface(SmoothScrollArea):
    def __init__(self, parent=None):
        super().__init__(parent=parent)
        self.view = QWidget(self)
        self.vBoxLayout = QVBoxLayout(self.view)
        self.setWidget(self.view)
        self.setWidgetResizable(True)

        # 将背景样式直接应用到主视图
        self.view.setStyleSheet('''
            QWidget#mainView {
                border-image: url(pic/background2.png) 0 0 0 0 stretch stretch;
            }
        ''')
        self.view.setObjectName("mainView")

        # 确保滚动区域和其他组件保持透明
        self.setStyleSheet('''
            QScrollArea {
                border: none;
                background: transparent;
            }
            QLabel {
                background: transparent;
                color: white;
            }
            QPushButton {
                border: none;
                background-color: transparent;
            }
        ''')

        plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'Arial']
        plt.rcParams['axes.unicode_minus'] = False
        plt.rcParams['axes.grid'] = True
        self.is_first_plot = True
        self.current_mode = 'drone'  # 当前图表模式，默认为无人机数量与总路径成本
        self.initUI()

    def initUI(self):
        self.vBoxLayout.setAlignment(Qt.AlignCenter)
        self.vBoxLayout.setSpacing(30)  # 增大控件间距

        # 添加顶部空白
        self.vBoxLayout.addSpacing(150)

        # 添加标题并设置字体
        title = QLabel("统计结果")
        title_font = QFont("楷体", 36, QFont.Bold)
        title.setFont(title_font)
        title.setAlignment(Qt.AlignCenter)
        self.vBoxLayout.addWidget(title)

        # 增加标题和图表之间的间距
        self.vBoxLayout.addSpacing(0)

        # 创建图表
        self.figure = plt.figure(figsize=(12, 6))  # 增大图表尺寸
        self.canvas = FigureCanvasQTAgg(self.figure)
        self.vBoxLayout.addWidget(self.canvas)

        # 增加图表和按钮之间的间距
        self.vBoxLayout.addSpacing(60)

        # 创建第一行按钮水平布局
        first_row_button_layout = QHBoxLayout()
        first_row_button_layout.setAlignment(Qt.AlignCenter)

        # 添加刷新按钮
        self.refresh_button = PrimaryPushButton('刷新数据')
        self.refresh_button.setFixedWidth(240)  # 加宽按钮
        self.refresh_button.setFixedHeight(50)  # 加高按钮
        self.refresh_button.setFont(QFont("Microsoft YaHei", 13))
        self.refresh_button.clicked.connect(self.update_plot)
        first_row_button_layout.addWidget(self.refresh_button)

        # 添加加载新数据按钮
        self.load_new_data_button = PrimaryPushButton('加载数据')
        self.load_new_data_button.setFixedWidth(240)
        self.load_new_data_button.setFixedHeight(50)
        self.load_new_data_button.setFont(QFont("Microsoft YaHei", 13))
        self.load_new_data_button.clicked.connect(self.load_new_data)
        first_row_button_layout.addWidget(self.load_new_data_button)

        # 添加清空数据按钮
        self.clear_data_button = PrimaryPushButton('清空数据')
        self.clear_data_button.setFixedWidth(240)
        self.clear_data_button.setFixedHeight(50)
        self.clear_data_button.setFont(QFont("Microsoft YaHei", 13))
        self.clear_data_button.clicked.connect(self.clear_plot)
        first_row_button_layout.addWidget(self.clear_data_button)

        self.vBoxLayout.addLayout(first_row_button_layout)

        # 创建第二行按钮水平布局
        second_row_button_layout = QHBoxLayout()
        second_row_button_layout.setAlignment(Qt.AlignCenter)

        # 添加显示无人机数量与路径成本的关系按钮
        self.switch_drone_button = PrimaryPushButton('显示无人机数量与路径成本的关系')
        # 增大按钮宽度以适应长文本
        self.switch_drone_button.setFixedWidth(320)
        self.switch_drone_button.setFixedHeight(50)
        self.switch_drone_button.setFont(QFont("Microsoft YaHei", 13))
        self.switch_drone_button.clicked.connect(lambda: self.switch_plot_mode('drone'))
        second_row_button_layout.addWidget(self.switch_drone_button)

        # 添加显示障碍物数量与路径成本按钮
        self.switch_obstacle_button = PrimaryPushButton('显示障碍物数量与路径成本的关系')
        # 增大按钮宽度以适应长文本
        self.switch_obstacle_button.setFixedWidth(320)
        self.switch_obstacle_button.setFixedHeight(50)
        self.switch_obstacle_button.setFont(QFont("Microsoft YaHei", 13))
        self.switch_obstacle_button.clicked.connect(lambda: self.switch_plot_mode('obstacle'))
        second_row_button_layout.addWidget(self.switch_obstacle_button)

        # 添加显示显示空域大小与路径成本的关系按钮
        self.switch_airspace_button = PrimaryPushButton('显示空域大小与路径成本的关系')
        # 增大按钮宽度以适应长文本
        self.switch_airspace_button.setFixedWidth(320)
        self.switch_airspace_button.setFixedHeight(50)
        self.switch_airspace_button.setFont(QFont("Microsoft YaHei", 13))
        self.switch_airspace_button.clicked.connect(lambda: self.switch_plot_mode('airspace'))
        second_row_button_layout.addWidget(self.switch_airspace_button)

        self.vBoxLayout.addLayout(second_row_button_layout)

        # 初始隐藏显示无人机数量与路径成本的关系按钮
        self.switch_drone_button.hide()

        # 添加底部弹性空间
        self.vBoxLayout.addStretch()

    def showEvent(self, event):
        super().showEvent(event)
        self.update_plot()

    def update_plot(self):
        try:
            if self.is_first_plot:
                self.figure.clear()
                self.is_first_plot = False

            current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
            file_path = os.path.join(current_dir, 'data', 'total_cost.txt')

            if not os.path.exists(file_path):
                raise Exception("未找到统计数据文件，请先进行路径规划")

            num_agents, total_costs = read_cost_data(file_path, self.current_mode)
            if not num_agents or not total_costs:
                raise Exception("没有找到有效的统计数据")

            if len(self.figure.axes) == 0:
                ax = self.figure.add_subplot(111)
            else:
                ax = self.figure.axes[0]

            ax.clear()  # 清除原图表内容

            # 设置图表样式
            color = COLOR_LIST[0]
            ax.scatter(num_agents, total_costs, s=150, color=color)  # 增大散点大小并设置颜色
            ax.plot(num_agents, total_costs, '--', alpha=0.6, color=color)  # 匹配散点颜色

            if self.current_mode == 'drone':
                ax.set_title('无人机数量与总路径成本关系', fontsize=16, pad=20)  # 增大标题字号
                ax.set_xlabel('无人机数量', fontsize=14, labelpad=10)  # 增大轴标签字号
            elif self.current_mode == 'obstacle':
                ax.set_title('障碍物数量与路径成本关系', fontsize=16, pad=20)
                ax.set_xlabel('障碍物数量', fontsize=14, labelpad=10)
            elif self.current_mode == 'airspace':
                ax.set_title('空域大小与路径成本关系', fontsize=16, pad=20)
                ax.set_xlabel('空域大小', fontsize=14, labelpad=10)

            ax.set_ylabel('总路径成本', fontsize=14, labelpad=10)
            ax.tick_params(labelsize=12)  # 增大刻度标签字号
            ax.grid(True, linestyle='--', alpha=0.3)

            for x, y in zip(num_agents, total_costs):
                ax.annotate(f'({x}, {int(y)})',
                            (x, y), textcoords="offset points",
                            xytext=(0, 10), ha='center')

            self.figure.tight_layout()
            self.canvas.draw()

            InfoBar.success(
                title='成功',
                content='统计图表已更新',
                orient=Qt.Horizontal,
                isClosable=True,
                position=InfoBarPosition.TOP,
                duration=2000,
                parent=self
            )
        except Exception as e:
            InfoBar.error(
                title='错误',
                content=str(e),
                orient=Qt.Horizontal,
                isClosable=True,
                position=InfoBarPosition.TOP,
                duration=4000,
                parent=self
            )

    def load_new_data(self):
        global current_color_index
        try:
            file_dialog = QFileDialog()
            file_path, _ = file_dialog.getOpenFileName(self, "选择数据文件", "", "文本文件 (*.txt)")

            if not file_path:
                return

            new_num_agents, new_total_costs = read_cost_data(file_path, self.current_mode)
            if not new_num_agents or not new_total_costs:
                raise Exception("所选文件中没有找到有效的统计数据")

            if len(self.figure.axes) == 0:
                ax = self.figure.add_subplot(111)
            else:
                ax = self.figure.axes[0]

            # 选择下一个颜色
            color = COLOR_LIST[current_color_index % len(COLOR_LIST)]
            current_color_index += 1

            # 不清除原图表，继续在同一图表上绘制新数据
            scatter = ax.scatter(new_num_agents, new_total_costs, s=150, color=color)
            line, = ax.plot(new_num_agents, new_total_costs, '--', alpha=0.6, color=color)

            for x, y in zip(new_num_agents, new_total_costs):
                ax.annotate(f'({x}, {int(y)})',
                            (x, y), textcoords="offset points",
                            xytext=(0, 10), ha='center')

            # 标注文件名，去除文件类型
            file_name = os.path.splitext(os.path.basename(file_path))[0]
            x_mid = (min(new_num_agents) + max(new_num_agents)) / 2
            y_mid = (min(new_total_costs) + max(new_total_costs)) / 2
            ax.annotate(file_name, (x_mid, y_mid), textcoords="offset points",
                        xytext=(0, 20), ha='center', color=color)

            self.figure.tight_layout()
            self.canvas.draw()

            InfoBar.success(
                title='成功',
                content='新统计数据已加载并显示',
                orient=Qt.Horizontal,
                isClosable=True,
                position=InfoBarPosition.TOP,
                duration=2000,
                parent=self
            )
        except Exception as e:
            InfoBar.error(
                title='错误',
                content=str(e),
                orient=Qt.Horizontal,
                isClosable=True,
                position=InfoBarPosition.TOP,
                duration=4000,
                parent=self
            )

    def switch_plot_mode(self, mode):
        self.current_mode = mode
        if mode == 'drone':
            self.switch_drone_button.hide()
            self.switch_obstacle_button.show()
            self.switch_airspace_button.show()
        elif mode == 'obstacle':
            self.switch_obstacle_button.hide()
            self.switch_drone_button.show()
            self.switch_airspace_button.show()
        elif mode == 'airspace':
            self.switch_airspace_button.hide()
            self.switch_drone_button.show()
            self.switch_obstacle_button.show()
        self.update_plot()

    def clear_plot(self):
        try:
            if len(self.figure.axes) > 0:
                ax = self.figure.axes[0]
                ax.clear()
                ax.set_xlabel('')
                ax.set_ylabel('')
                ax.set_title('')
                self.figure.tight_layout()
                self.canvas.draw()
            InfoBar.success(
                title='成功',
                content='图表数据已清空',
                orient=Qt.Horizontal,
                isClosable=True,
                position=InfoBarPosition.TOP,
                duration=2000,
                parent=self
            )
        except Exception as e:
            InfoBar.error(
                title='错误',
                content=str(e),
                orient=Qt.Horizontal,
                isClosable=True,
                position=InfoBarPosition.TOP,
                duration=4000,
                parent=self
            )
